from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 加载模型（这里用 GPT-2，你可以换成 LLaMA, Falcon 等）
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 输入提示文本
prompt = "从前山里面有个和尚 挑水吃"

# 分词并转成 PyTorch 张量
inputs = tokenizer(prompt, return_tensors="pt")  #pt就是 PyTorch

# 生成文本
outputs = model.generate(
    **inputs,
    max_length=500,       # 生成的总token数（包含输入部分）
    temperature=0.7,     # 控制多样性，越高越随机
    top_p=0.9,           # nucleus sampling
    do_sample=True,      # 启用采样（否则是贪心搜索） 固定、可控输出 → do_sample=False  自然、多样、创意输出 → do_sample=True
    repetition_penalty=1.2,  # 惩罚重复
    pad_token_id=tokenizer.eos_token_id  # 防止报错
)

# 解码并打印
result = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(result)
#对话系统（基于历史消息生成下一条回复）
#代码生成（如 CodeLLaMA、StarCoder）
#故事创作（长文本续写）
#多轮问答（配合 prompt 工程）
#基本上上次小故事 扩展什么的